Creating a Dockerfile#69
Merged
Merged
Conversation
…tandardize the edia data dir to be the same as pdb dir
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdds GPU-capable Docker containerization (Dockerfile, docker-compose, entrypoint, .dockerignore) and refactors EDIA handling from per‑PDB CSVs/ Changes
Sequence Diagram(s)sequenceDiagram
participant Dev as Developer
participant Docker as Docker Build
participant Builder as Builder Stage (CUDA)
participant HF as HuggingFace
participant Runtime as Runtime Container
participant Entrypoint as /app/entrypoint.sh
Dev->>Docker: docker build / docker-compose build
Docker->>Builder: run build steps (apt, python, uv sync, venv)
Builder->>HF: pre-download esm3-open into HF cache
Builder->>Builder: compile bytecode, run import/CUDA checks
Docker->>Runtime: assemble runtime image (venv, src, cache)
Dev->>Runtime: docker-compose run (mount volumes, env vars)
Runtime->>Entrypoint: execute command (train / inference / generate-esm / python)
Entrypoint->>Runtime: launch corresponding python script with WATERFLOW_* paths
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds containerization (Dockerfile, docker-compose, entrypoint) to run training/inference/embedding generation in a CUDA 12.6 environment, and simplifies EDIA-based water filtering by switching from a separate EDIA CSV directory to a JSON file colocated with each PDB-REDO structure.
Changes:
- Add a multi-stage CUDA Docker image, docker-compose configuration, and a unified entrypoint CLI.
- Switch EDIA loading from per-PDB CSVs in an external
edia_dirto a colocated JSON file; remove--edia_dirplumbing from train/inference and docs. - Update README to reflect the new EDIA JSON location/format.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
src/dataset.py |
Reworks EDIA loader to read JSON and removes edia_dir dependency; updates preprocessing to derive JSON path from PDB path. |
scripts/train.py |
Removes --edia_dir CLI option and stops passing edia_dir into dataset config. |
scripts/inference.py |
Stops extracting/passing edia_dir from training config into dataset filter config. |
docker/entrypoint.sh |
Adds a command dispatcher for train, inference, and generate-esm inside the container. |
docker-compose.yml |
Adds compose setup for GPU-enabled runs and a dev profile. |
README.md |
Removes --edia_dir docs and updates EDIA description to JSON colocated with PDB. |
Dockerfile |
Adds multi-stage build using uv, pre-downloads ESM3, and sets up runtime image/entrypoint. |
.dockerignore |
Excludes local/dev artifacts and large files from Docker build context. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
Dockerfile (2)
29-29: Pin theuvimage instead oflatest.Using
ghcr.io/astral-sh/uv:latestmakes the build non-reproducible; an upstream tag move can change tooling behavior or break the image without any repo diff. Please pin the version or digest that CI has validated.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@Dockerfile` at line 29, The Dockerfile currently copies the uv binary from an unpinned image reference "ghcr.io/astral-sh/uv:latest" which makes builds non-reproducible; update the COPY statement to use a specific, pinned tag or digest for the "ghcr.io/astral-sh/uv" image (e.g., a validated semver tag or `@sha256`:... digest) so CI-validated image is used consistently—replace the reference in the COPY line that mentions "ghcr.io/astral-sh/uv:latest" with the chosen pinned tag/digest.
49-55: Make the ESM cache warm-up opt-in.Every build downloads
esm3-open, even when the image is only used for GVP/SLAE workflows. That adds a large, network-dependent layer to all builds and makes otherwise valid non-ESM builds slower and more brittle. A build arg or separate target would keep the default image lighter.One way to gate the warm-up
+# Optional: prewarm the ESM cache for images that need embedding generation. +ARG PREWARM_ESM=0 ENV HF_HOME=/app/.cache/huggingface -RUN . .venv/bin/activate && python -c "\ -from esm.models.esm3 import ESM3; \ -ESM3.from_pretrained('esm3-open'); \ -print('ESM3 model downloaded successfully')" +RUN if [ "$PREWARM_ESM" = "1" ]; then \ + . .venv/bin/activate && python -c "\ +from esm.models.esm3 import ESM3; \ +ESM3.from_pretrained('esm3-open'); \ +print('ESM3 model downloaded successfully')"; \ + fi🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@Dockerfile` around lines 49 - 55, Make the ESM3 cache warm-up opt-in by gating the RUN snippet that calls ESM3.from_pretrained('esm3-open') behind a build arg or separate target: add an ARG like PREWARM_ESM (default "no") and wrap the existing RUN ". .venv/bin/activate && python -c 'from esm.models.esm3 import ESM3; ESM3.from_pretrained(\"esm3-open\"); ...'" in a shell conditional (if [ "$PREWARM_ESM" = "yes" ]; then ... fi) so the HF_HOME cache setup only runs when PREWARM_ESM=yes, or alternatively move that RUN into its own Docker target (e.g., "esm-warmup") and only invoke that target for builds that need the predownloaded model.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@docker/entrypoint.sh`:
- Around line 93-95: The bash) case in docker/entrypoint.sh currently calls exec
/bin/bash without forwarding arguments, so options like -lc and the command are
ignored; update that branch (the bash) case that invokes exec /bin/bash) to
forward the original container arguments (e.g., exec /bin/bash "$@" or otherwise
pass "${@}") so any flags and commands provided to docker run (like -lc '...')
are honored.
- Around line 72-77: The inference branch in the entrypoint script fails to
forward the WATERFLOW_OUTPUT_DIR env var to the python invocation, so
scripts/inference.py still errors waiting for --output_dir; update the inference
case (the exec python /app/scripts/inference.py command) to include --output_dir
"${WATERFLOW_OUTPUT_DIR}" (properly quoted) before the "$@" arguments so the
container passes the defined env var into the script.
In `@Dockerfile`:
- Around line 97-126: The Dockerfile currently leaves the final image running as
root (after COPY/ENTRYPOINT and RUN mkdir steps), so create a non-root user and
make the runtime switch: add a user/group (e.g., appuser, UID 1000), chown the
application directories and the mounted data directories (/app,
/app/entrypoint.sh and the created /data/* paths) to that user, ensure
executable bit on /app/entrypoint.sh remains set, and add a USER appuser
directive before ENTRYPOINT so the container runs as the non-root user; update
any ENV paths ownership (VIRTUAL_ENV, HF_HOME) as needed so the non-root user
can write to them.
In `@src/dataset.py`:
- Around line 462-464: The code unconditionally sets edia_lookup[key] =
float(entry.get("EDIAm", 0.0)) which converts missing or null EDIAm into 0.0 or
raises a TypeError; instead check entry.get("EDIAm") is present and numeric
before adding the lookup. In the block that builds key = (chain_id, res_id,
ins_code) only add edia_lookup[key] when entry["EDIAm"] exists and can be
converted to float (e.g., not None and is a number/string parsable as float),
otherwise skip adding the key so apply_threshold_filter() can conservatively
treat missing values. Ensure you reference and guard conversions around "EDIAm"
when populating edia_lookup.
---
Nitpick comments:
In `@Dockerfile`:
- Line 29: The Dockerfile currently copies the uv binary from an unpinned image
reference "ghcr.io/astral-sh/uv:latest" which makes builds non-reproducible;
update the COPY statement to use a specific, pinned tag or digest for the
"ghcr.io/astral-sh/uv" image (e.g., a validated semver tag or `@sha256`:...
digest) so CI-validated image is used consistently—replace the reference in the
COPY line that mentions "ghcr.io/astral-sh/uv:latest" with the chosen pinned
tag/digest.
- Around line 49-55: Make the ESM3 cache warm-up opt-in by gating the RUN
snippet that calls ESM3.from_pretrained('esm3-open') behind a build arg or
separate target: add an ARG like PREWARM_ESM (default "no") and wrap the
existing RUN ". .venv/bin/activate && python -c 'from esm.models.esm3 import
ESM3; ESM3.from_pretrained(\"esm3-open\"); ...'" in a shell conditional (if [
"$PREWARM_ESM" = "yes" ]; then ... fi) so the HF_HOME cache setup only runs when
PREWARM_ESM=yes, or alternatively move that RUN into its own Docker target
(e.g., "esm-warmup") and only invoke that target for builds that need the
predownloaded model.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a3af767f-641e-4eca-a339-e4dddea5f31a
📒 Files selected for processing (8)
.dockerignoreDockerfileREADME.mddocker-compose.ymldocker/entrypoint.shscripts/inference.pyscripts/train.pysrc/dataset.py
💤 Files with no reviewable changes (2)
- scripts/inference.py
- scripts/train.py
There was a problem hiding this comment.
♻️ Duplicate comments (1)
src/dataset.py (1)
463-465:⚠️ Potential issue | 🟠 MajorSkip missing/null
EDIAmvalues instead of turning them into failures.The current implementation at line 465 has two issues:
- Missing
EDIAmdefaults to0.0, which will incorrectly fail the EDIA threshold filter- If
EDIAmis explicitlynullin JSON,float(None)raisesTypeErrorOnly add the key when
EDIAmis actually present and numeric.
,Suggested fix
# Build the lookup key and extract the EDIAm score key = (chain_id, res_id, ins_code) - edia_lookup[key] = float(entry.get("EDIAm", 0.0)) + ediam = entry.get("EDIAm") + if ediam is not None: + edia_lookup[key] = float(ediam)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/dataset.py` around lines 463 - 465, The code unconditionally casts entry.get("EDIAm", 0.0) to float and stores it in edia_lookup causing false zeros and TypeError for nulls; change this to check entry for a present, non-null, numeric EDIAm before adding to edia_lookup: read value = entry.get("EDIAm"), verify value is not None and is an int/float or convertible to float (or use isinstance(value, (int, float))), then set edia_lookup[key] = float(value) only when valid; leave out the key when EDIAm is missing or not numeric. Use the existing variables key, entry, and edia_lookup to locate the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Duplicate comments:
In `@src/dataset.py`:
- Around line 463-465: The code unconditionally casts entry.get("EDIAm", 0.0) to
float and stores it in edia_lookup causing false zeros and TypeError for nulls;
change this to check entry for a present, non-null, numeric EDIAm before adding
to edia_lookup: read value = entry.get("EDIAm"), verify value is not None and is
an int/float or convertible to float (or use isinstance(value, (int, float))),
then set edia_lookup[key] = float(value) only when valid; leave out the key when
EDIAm is missing or not numeric. Use the existing variables key, entry, and
edia_lookup to locate the change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b18b0ab3-0cc6-4b76-8b7c-8a479fbddb72
📒 Files selected for processing (2)
src/dataset.pytests/test_dataset.py
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
src/dataset.py (1)
467-470: Returning{}vsNonecreates distinct behaviors — document or unify.
None→ file missing or parse error → raisesValueErrorin_preprocess_one{}→ file valid but no water entries → no error, empty lookup passed to filterThis distinction seems intentional (empty file is valid, missing file is not), but the docstring only mentions
Nonefor errors. Consider documenting{}in the docstring for clarity:📝 Docstring update
Returns: Dictionary mapping (chain_id, res_id, ins_code) -> EDIA score for waters, - or None if file not found or error + empty dict if file contains no water entries, + or None if file not found or parse error🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/dataset.py` around lines 467 - 470, The function that returns edia_lookup currently returns {} for a valid-but-empty EDIA file and None for missing/parse errors, causing different behaviors in callers like _preprocess_one; update the docstring of the function that returns edia_lookup to explicitly state that None indicates a missing/parse error while {} indicates a valid file with no water entries (and note how _preprocess_one will treat each), or alternatively make the behavior uniform by returning None for errors only and keeping {} for empty valid files—ensure the docstring and _preprocess_one handling are consistent with whichever choice you make.Dockerfile (1)
87-89:TORCH_CUDA_ARCH_LIST="9.0"restricts to H100 GPUs only.This architecture list targets only Hopper (H100). If users need to run on older GPUs (A100=8.0, V100=7.0, etc.), they'll need to rebuild the image or override this at runtime. If H100-only is intentional, consider adding a comment noting this restriction.
💡 Optional: Document or expand architecture support
# CUDA configuration for H100 GPUs ENV CUDA_HOME=/usr/local/cuda +# Note: "9.0" targets H100 only. For broader GPU support, use e.g. "7.0;7.5;8.0;8.6;9.0" ENV TORCH_CUDA_ARCH_LIST="9.0"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@Dockerfile` around lines 87 - 89, The Dockerfile currently hard-sets TORCH_CUDA_ARCH_LIST="9.0", which restricts builds to H100 GPUs; update the Dockerfile to either broaden supported architectures (e.g., include "8.0;9.0" or other required SMs) or leave TORCH_CUDA_ARCH_LIST unset/empty so users can build for their target GPUs at runtime, and if H100-only is intentional add a clear comment next to TORCH_CUDA_ARCH_LIST explaining the restriction; reference the environment variable TORCH_CUDA_ARCH_LIST and the existing CUDA_HOME assignment when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/dataset.py`:
- Around line 456-457: The code currently defaults pdb_info.get("seqNum", 0)
into res_id which can silently produce res_id==0 and cause collisions; update
the handling in the block that assigns chain_id/res_id (using variables
pdb_info, chain_id, res_id) to first check that "seqNum" exists and is not
null/empty, attempt safe int conversion (catch ValueError), and if
missing/invalid either log a warning including the offending pdb_info (or its
identifying fields) and skip processing that entry rather than defaulting to 0;
ensure downstream logic that expects res_id only runs when a valid res_id was
obtained.
---
Nitpick comments:
In `@Dockerfile`:
- Around line 87-89: The Dockerfile currently hard-sets
TORCH_CUDA_ARCH_LIST="9.0", which restricts builds to H100 GPUs; update the
Dockerfile to either broaden supported architectures (e.g., include "8.0;9.0" or
other required SMs) or leave TORCH_CUDA_ARCH_LIST unset/empty so users can build
for their target GPUs at runtime, and if H100-only is intentional add a clear
comment next to TORCH_CUDA_ARCH_LIST explaining the restriction; reference the
environment variable TORCH_CUDA_ARCH_LIST and the existing CUDA_HOME assignment
when making the change.
In `@src/dataset.py`:
- Around line 467-470: The function that returns edia_lookup currently returns
{} for a valid-but-empty EDIA file and None for missing/parse errors, causing
different behaviors in callers like _preprocess_one; update the docstring of the
function that returns edia_lookup to explicitly state that None indicates a
missing/parse error while {} indicates a valid file with no water entries (and
note how _preprocess_one will treat each), or alternatively make the behavior
uniform by returning None for errors only and keeping {} for empty valid
files—ensure the docstring and _preprocess_one handling are consistent with
whichever choice you make.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a65bfd6d-de5f-4783-a340-25848bad6458
📒 Files selected for processing (3)
Dockerfilesrc/dataset.pytests/test_dataset.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/test_dataset.py
marcuscollins
approved these changes
Mar 26, 2026
marcuscollins
left a comment
marcuscollins
left a comment
Collaborator
There was a problem hiding this comment.
a few last comments to create issues and think about for the future, otherwise lgtm.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Open
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR towards:
load_ediamethod in dataset.py to read inpdb_dir/<pdb_id>_final.jsoncolocated with the pdb file instead of a csv in a separate edia directory to simplify the filtering when processing. Thejsonfile is available to download from PDB-REDO alongside the pdb file. The reason for adding this change in this PR is to ensure I do not have to specify another path variable (edia_dir).Summary by CodeRabbit
New Features
Documentation
Chores
Tests
Bug Fixes